# -*- coding: utf-8 -*-
"""
裁剪
"""

import sys
import argparse
import os
import random
import cv2

import voc_xml
from voc_xml import CreateXML


def confine(value, v_min, v_max):
    """
    值的边界限制
    Args:
        value:输入值
        v_min,v_max:最大最小边界
    return:
        value:限制值
    """
    value = v_min if value < v_min else value
    value = v_max if value > v_max else value
    return value


def calc_rect_area(rect):
    """
    计算矩形框面积
    Args:
        rect:矩形框 [xmin,ymin,xmax,ymax]
    return:
        dst:矩形框面积
    """
    return (rect[2] - rect[0] + 0.001) * (rect[3] - rect[1] + 0.001)


# rect : [left, top, right, bottom]
def calc_iou(rect1, rect2):
    """
    计算两个矩形框的交并比
    Args:
        rect1,rect2:两个矩形框
    return:
        iou:交并比
    """
    rect_intersection = (max(rect1[0], rect2[0]), max(rect1[1], rect2[1]),
                         min(rect1[2], rect2[2]), min(rect1[3], rect2[3]))
    intersection_w = rect_intersection[2] - rect_intersection[0] + 0.001
    intersection_h = rect_intersection[3] - rect_intersection[1] + 0.001

    iou = 0
    if intersection_w > 0 and intersection_h > 0:
        union_area = calc_rect_area(rect1) + calc_rect_area(rect2) - \
                     intersection_w * intersection_h
        iou = intersection_w * intersection_h / union_area
    return iou


def crop_img(src, left, top, crop_w, crop_h):
    """
    裁剪图像
    Args:
        src: 源图像
        top_left,top_right:裁剪图像左上角坐标
        crop_w,crop_h：裁剪图像宽高
    return：
        crop_img:裁剪后的图像
        None:裁剪尺寸错误
    """
    rows, cols, channels = src.shape
    row_min, col_min = int(top), int(left)
    row_max, col_max = int(row_min + crop_h), int(col_min + crop_w)
    if row_max > rows or col_max > cols:
        print("crop size err: src->%dx%d,crop->top_left(%d,%d) %dx%d" % (
            cols, rows, col_min, row_min, int(crop_w), int(crop_h)))
        return None
    crop_img = src[row_min:row_max, col_min:col_max]
    return crop_img


def crop_xy(x, y, top_left_x, top_left_y, crop_w, crop_h):
    """
    坐标平移变换
    Args:
        x,y:待变换坐标
        top_left_x,top_left_y:裁剪图像左上角坐标
        crop_w,crop_h:裁剪部分图像宽高
    return:
        crop_x,crop_y
    """
    crop_x = int(x - top_left_x)
    crop_y = int(y - top_left_y)
    crop_x = confine(crop_x, 0, crop_w - 1)
    crop_y = confine(crop_y, 0, crop_h - 1)
    return crop_x, crop_y


def crop_box(box, top_left_x, top_left_y, crop_w, crop_h, iou_thr=0.5):
    """
    目标框坐标平移变换
    Args:
        box:目标框坐标[xmin,ymin,xmax,ymax]
        top_left_x,top_left_y:裁剪图像左上角坐标
        crop_w,crop_h:裁剪部分图像宽高
        iou_thr: iou阈值,去除裁剪后过小目标
    return:
        crop_box:平移变换结果[xmin,ymin,xmax,ymax]
    """
    xmin, ymin = crop_xy(box[0], box[1], top_left_x, top_left_y, crop_w, crop_h)
    xmax, ymax = crop_xy(box[2], box[3], top_left_x, top_left_y, crop_w, crop_h)
    croped_box = [xmin, ymin, xmax, ymax]
    if calc_iou([0, 0, box[2] - box[0], box[3] - box[1]],
                [0, 0, xmax - xmin, ymax - ymin]) < iou_thr:
        croped_box = [0, 0, 0, 0]
    return croped_box


def crop_xml(crop_img_name, xml_tree, top_left_x,
             top_left_y, crop_w, crop_h, iou_thr=0.5):
    """
    xml目标框裁剪变换
    Args:
        crop_img_name:裁剪图片命名
        xml_tree：待crop的xml ET.parse()
        top_left_x,top_left_y: 裁剪图像左上角坐标
        crop_w,crop_h: 裁剪图像宽高
        iou_thr: iou阈值
    return:
        createdxml : 创建的xml CreateXML对象         
    """
    root = xml_tree.getroot()
    size = root.find('size')
    depth = int(size.find('depth').text)
    createdxml = CreateXML(crop_img_name, int(crop_w), int(crop_h), depth)
    for obj in root.iter('object'):
        obj_name = obj.find('name').text
        xml_box = obj.find('bndbox')
        xmin = int(xml_box.find('xmin').text)
        ymin = int(xml_box.find('ymin').text)
        xmax = int(xml_box.find('xmax').text)
        ymax = int(xml_box.find('ymax').text)
        box = crop_box([xmin, ymin, xmax, ymax],
                       top_left_x, top_left_y, crop_w, crop_h, iou_thr)
        if (box[0] >= box[2]) or (box[1] >= box[3]):
            continue
        createdxml.add_object_node(obj_name, box[0], box[1], box[2], box[3])
    return createdxml


def crop_img_xml(img, xml_tree, crop_img_name, left, top,
                 crop_w, crop_h, iou_thr):
    """
    裁剪图像和xml目标框
    Args:
        img：源图像
        crop_img_name:裁剪图片命名
        xml_tree：待crop的xml ET.parse()
        top_left_x,top_left_y: 裁剪图像左上角坐标
        crop_w,crop_h: 裁剪图像宽高
        iou_thr: iou阈值
    return:
        croped_img,croped_xml : 裁剪完成的图像和xml文件
        None:裁剪尺寸错误
    """
    croped_img = crop_img(img, left, top, crop_w, crop_h)
    if croped_img is None:
        return None, None
    croped_xml = crop_xml(crop_img_name, xml_tree, left, top,
                          crop_w, crop_h, iou_thr)
    return croped_img, croped_xml


def batch_crop_dir(imgs_dir, xmls_dir, imgs_save_dir,
                   xmls_save_dir, img_suffix, name_suffix,
                   crop_type='RANDOM_CROP', crop_n=1, crop_size=(0, 0),
                   scale_w=1.0, scale_h=1.0, random_wh=False, iou_thr=0.5):
    """
    batch crop images in directory
    :param imgs_dir: JPEGImages
    :param xmls_dir: Annotations
    :param imgs_save_dir: crop image save dir
    :param xmls_save_dir: crop xml save dir
    :param img_suffix: .jpg .png .bmp
    :param name_suffix: to identify crop method
    :param crop_type: crop method
    :param crop_n: to generate crop_n images per original image
    :param crop_size: crop size
    :param scale_w: ratio to width. only used when crop_size=(0, 0)
    :param scale_h: ratio to height. only used when crop_size=(0, 0)
    :param random_wh: random crop flag
    :param iou_thr: iou threshold to decide if an object should be kept
    :return: None
    """
    for root, dirs, files in os.walk(xmls_dir):
        for xml_name in files:
            xml_file = os.path.join(xmls_dir, xml_name)
            img_file = None
            for suffix in img_suffix:
                if os.path.exists(os.path.join(
                        imgs_dir, xml_name.split('.')[0]+suffix)):
                    img_file = os.path.join(
                        imgs_dir, xml_name.split('.')[0]+suffix)
                    break
            if img_file is None:
                print("there has no image for ", xml_name)
                continue
            img = cv2.imread(img_file)
            img_h, img_w, img_channels = img.shape

            if crop_type == 'CENTER_CROP':
                crop_n = 1
            elif crop_type == 'FIVE_CROP':
                crop_n = 5

            crop_left, crop_top = 0, 0
            for i in range(crop_n):
                crop_w, crop_h = crop_size
                if crop_size == (0, 0) and not random_wh:
                    crop_w = int(img_w * scale_w)
                    crop_h = int(img_h * scale_h)
                elif random_wh:
                    crop_w = int(img_w * (scale_w + random.random() * (1 - scale_w)))
                    crop_h = int(img_h * (scale_h + random.random() * (1 - scale_h)))

                if crop_type == 'RANDOM_CROP':
                    crop_left, crop_top = random.randint(0, img_w - crop_w - 1),\
                                          random.randint(0, img_h - crop_h - 1)
                elif crop_type == 'CENTER_CROP':
                    crop_left, crop_top = int((img_w - crop_w) / 2),\
                                          int((img_h - crop_h) / 2)
                elif crop_type == 'FIVE_CROP':
                    if i == 0:
                        crop_left, crop_top = 0, 0
                    elif i == 1:
                        crop_left, crop_top = img_w - crop_w - 1, 0
                    elif i == 2:
                        crop_left, crop_top = 0, img_h - crop_h - 1
                    elif i == 3:
                        crop_left, crop_top = img_w - crop_w - 1,\
                                              img_h - crop_h - 1
                    else:
                        crop_left, crop_top = int((img_w - crop_w) / 2),\
                                              int((img_h - crop_h) / 2)
                else:
                    print('crop type wrong! expect [RANDOM_CROP, CENTER_CROP, FIVE_CROP]')
                cropped_img_name = "{}_{}_left{}_top{}_{}x{}.{}".format(
                    xml_name.split('.')[0], name_suffix, crop_left, crop_top,
                    crop_w, crop_h, img_file.split('.')[-1])

                img_crop, xml_crop = crop_img_xml(img, voc_xml.get_xml_tree(xml_file),
                                                  cropped_img_name, crop_left, crop_top,
                                                  crop_w, crop_h, iou_thr)
                if img_crop is not None:
                    cv2.imwrite(os.path.join(imgs_save_dir, cropped_img_name), img_crop)
                    xml_crop.save_xml(xmls_save_dir, cropped_img_name.split('.')[0]+'.xml')


def start_crop(root_dir, save_dir):
    img_dir = os.path.join(root_dir, "JPEGImages")
    xml_dir = os.path.join(root_dir, "Annotations")
    crop_img_save_dir = os.path.join(save_dir, "JPEGImages")
    crop_xml_save_dir = os.path.join(save_dir, "Annotations")
    if not os.path.exists(crop_img_save_dir):
        os.makedirs(crop_img_save_dir)
    if not os.path.exists(crop_xml_save_dir):
        os.makedirs(crop_xml_save_dir)

    img_suffix = ['.jpg', '.png', '.bmp']

    # add to file name to identify crop methods
    name_suffix = 'fivecrop'

    # ['RANDOM_CROP','CENTER_CROP','FIVE_CROP']
    crop_type = 'FIVE_CROP'

    # generate 'crop_n' cropped images every image
    crop_n = 5

    cropped_size = (300, 300)

    # crop according to fraction of width and height
    # only used when 'cropped_size' is not defined.
    scale_w = 0.5
    scale_h = 0.7

    # random crop flag
    random_wh = False

    # iou threshold to decide
    # if an object should be kept
    iou_thr = 0.5
    batch_crop_dir(img_dir, xml_dir, crop_img_save_dir,
                   crop_xml_save_dir, img_suffix, name_suffix,
                   crop_type, crop_n, cropped_size, scale_w,
                   scale_h, random_wh, iou_thr)


if __name__ == '__main__':
    try:
        ROOT_DIR = sys.argv[1]
        SAVE_DIR = sys.argv[2]
    except:
        PARSER = argparse.ArgumentParser(description="arguments")
        PARSER.add_argument("--root_dir", type=str,
                            default="/home/chujie/PycharmProjects/data_tool/data/voc2012",
                            help="VOC dataset root directory")
        PARSER.add_argument("--save_dir", type=str,
                            default=os.path.join(os.getenv("HOME"),
                                                 "PycharmProjects/data_tool/data/crop_voc"),
                            help="result save dir")
        ARGS = PARSER.parse_args()
        ROOT_DIR = ARGS.root_dir
        SAVE_DIR = ARGS.save_dir
    start_crop(ROOT_DIR, SAVE_DIR)
